import sqlparse
from dataclasses import dataclass
from typing import List
import re

@dataclass
class Column:
    name: str
    data_type: str
    comment: str = ""
    is_primary: bool = False
    is_foreign: bool = False
    reference_table: str = None
    reference_column: str = None

@dataclass
class Table:
    name: str
    columns: List[Column]
    comment: str = ""

class SQLParser:
    def parse(self, sql_text: str) -> List[Table]:
        tables = []
        # 移除单行SQL注释，但保留COMMENT内容
        sql_text = re.sub(r'--[^\n]*(?=\n|$)', '', sql_text)
        statements = sqlparse.parse(sql_text)
        
        for statement in statements:
            if not statement.get_type() == 'CREATE':
                continue
            
            # 获取表名和表注释
            table_name = None
            table_comment = ""
            columns = []
            
            # 提取表注释
            comment_match = re.search(r"COMMENT\s*=\s*'([^']*)'", str(statement))
            if comment_match:
                table_comment = comment_match.group(1)
            
            # 查找CREATE TABLE语句中的表名
            create_table_match = re.search(r"CREATE\s+TABLE\s+[`\"]?([^`\"\s]+)[`\"]?", str(statement), re.IGNORECASE)
            if create_table_match:
                table_name = create_table_match.group(1)
            
            # 查找列定义
            for token in statement.tokens:
                if isinstance(token, sqlparse.sql.Parenthesis):
                    columns = self._parse_columns(token.value)
            
            if table_name and columns:
                tables.append(Table(name=table_name, columns=columns, comment=table_comment))
        
        return tables
    
    def _parse_columns(self, columns_text: str) -> List[Column]:
        columns = []
        columns_text = columns_text.strip('()')
        
        # 分割列定义
        in_parentheses = 0
        current = []
        column_defs = []
        
        for char in columns_text:
            if char == '(' and in_parentheses >= 0:
                in_parentheses += 1
            elif char == ')' and in_parentheses > 0:
                in_parentheses -= 1
            elif char == ',' and in_parentheses == 0:
                column_defs.append(''.join(current).strip())
                current = []
                continue
            current.append(char)
        
        if current:
            column_defs.append(''.join(current).strip())

        # 处理每个列定义
        for col_def in column_defs:
            # 跳过外键和约束定义
            if col_def.upper().startswith(('PRIMARY KEY', 'FOREIGN KEY', 'CONSTRAINT', 'KEY')):
                continue
            
            # 提取列注释
            comment_match = re.search(r"COMMENT\s+'([^']*)'", col_def)
            comment = comment_match.group(1) if comment_match else ""
            
            # 移除COMMENT部分以便解析其他属性
            col_def = re.sub(r"COMMENT\s+'[^']*'", "", col_def)
            
            # 提取列名和数据类型
            name_type_match = re.match(r"`?([^`\s]+)`?\s+([^\s,]+)", col_def)
            if not name_type_match:
                continue
                
            name = name_type_match.group(1)
            data_type = name_type_match.group(2).upper()
            
            # 检查是否为主键
            is_primary = ('PRIMARY KEY' in col_def.upper() or 
                         'AUTO_INCREMENT' in col_def.upper())
            
            # 检查是否为外键 - 只检查明确的外键定义
            is_foreign = ('FOREIGN KEY' in col_def.upper() or 
                         'REFERENCES' in col_def.upper())
            
            column = Column(
                name=name,
                data_type=data_type,
                comment=comment,
                is_primary=is_primary,
                is_foreign=is_foreign
            )
            columns.append(column)
            
        return columns
    
    def _process_foreign_keys(self, tables: List[Table]):
        # 创建表名到表对象的映射
        table_map = {table.name: table for table in tables}
        
        # 遍历所有表和列，查找外键定义
        for table in tables:
            for column in table.columns:
                # 检查列名是否以_id结尾，这通常表示外键
                if column.name.lower().endswith('_id'):
                    # 推测引用表名（去掉_id后缀）
                    ref_table_name = column.name[:-3]
                    if ref_table_name in table_map:
                        column.is_foreign = True
                        column.reference_table = ref_table_name
                        column.reference_column = 'id'